# -*- coding: UTF-8 -*-
# python 连接kafka 消费数据

import json
import sys
import re
from importlib import import_module

import numpy
import pymysql
import pymysql.cursors
import torch
import torch.nn as nn
from kafka import KafkaConsumer

from utils.api_kafka.producer import Producer


class Consumer(object):
    #  GroupID='TestGroup001', ClientId="Test666",
    def __init__(self, model_name='bert', KafkaServerList=['127.0.0.1:9092'], GroupID='TestGroup0', ClientId="Test", Topics=['Test', ]):
        """
                :param KafkaServerList: kafka服务器IP:PORT 列表
                :param GroupID: 消费者组ID
                :param ClientId: 消费者名称
                :param Topic: 主题

        """
        self.model_name = model_name
        self._kwargs = {
            "bootstrap_servers": KafkaServerList,
            "client_id": ClientId,
            "group_id": GroupID,
            "enable_auto_commit": False,
            "auto_offset_reset": "earliest",
            "key_deserializer": lambda m: json.loads(m.decode('utf-8')),
            "value_deserializer": lambda m: json.loads(m.decode('utf-8')),
        }
        self.conn, self.cursor = self.connect2db('root', 'mzy123', 'jkw_demo', "10.10.149.20", 3306)
        try:
            self._consumer = KafkaConsumer(**self._kwargs)
            self._consumer.subscribe(topics=(Topics))
        except Exception as err:
            print("Consumer init failed, %s" % err)

    def consumeMsg(self, p):
        dataset = 'data/' + 'total2'
        x = import_module('models.' + self.model_name)
        config = (x.Config(dataset))
        model = x.Model(config).to(config.device)
        model.load_state_dict(torch.load(config.save_path, map_location=config.device))
        model.eval()
        cnt = 0
        last = ''
        try:
            while True:
                # print(cnt)
                #data = self._consumer.poll(timeout_ms=2, max_records=100)  # 拉取消息，字典类型
                data = self._consumer.poll(100)
                if data:
                    for key in data:
                        consumerrecord = data.get(key)[0]  # 返回的是ConsumerRecord对象，可以通过字典的形式获取内容。
                        if consumerrecord != None:
                            # 消息消费逻辑
                            message = {
                                "Topic": consumerrecord.topic,
                                "Partition": consumerrecord.partition,
                                "Offset": consumerrecord.offset,
                                "Key": consumerrecord.key,
                                "Value": consumerrecord.value
                            }
                            # print(consumerrecord.value['content'])
                            # print(consumerrecord.value)
                            str = self._clean(consumerrecord.value['content'])
                            if str == last:
                                self._consumer.commit()
                                continue
                            else: 
                                last = str
                            print(str)
                            # 分类结果为军事关心则输出到kafka
                            if self._deal_with(str, model):
                                print('guanzhu: 1')
                                print('分类结果: 军事关注, 输出至kafka')
                                p.sendMessage(message['Value'])
                            # 分类结果为军事非关系则输出到mysql
                            else:
                                print('guanzhu: 0')
                                print('分类结果: 军事非关注, 输出至mysql')
                                self._insertdatas(message['Value'])
                            cnt += 1
                            # 消费逻辑执行完毕后在提交偏移量
                            self._consumer.commit()
                        else:
                            print("%s consumerrecord is None." % key)
        except Exception as err:
            print(err)

    def _clean(self, message):
        res = re.sub(r'\^A', '', message)
        return res

    def _deal_with(self, message, model):
        softmax = nn.Softmax(dim=0)
        with torch.no_grad():
            outputs = model(message)
            for output in outputs:
                output = softmax(output)
                print(output)
                if output[0] - 0.2 > output[1]:
                    return 0
                else:
                    return 1
            
    def _insertdatas(self, data):
        # print(data)
        line = []
        line.append(int(float(data['id']))) # int函数不能接收str,需要float中转
        line.append(data['title'])
        line.append(data['content'])
        line.append(data['site'])
        line.append(data['url'])
        line.append(data['source'])
        line.append(data['time']) # publish time发送时间        
        line.append(data['news_author'])
        line.append(data['mainNavigation'])
        line.append(data['subNavigation'])
        line.append(data['inserTime']) # collect_time采集时间
        line.append(0) # pred_label_gz
        line.append(int(float(data['area']))) # pred_label_area
        try:
            self.conn.ping()
        except Exception as e:
            print("update scores mysql connection exception: %s"%e)
        
        # pred_label_area的默认值是0，但总是重写为junshi_guanzhu中area字段的值
        sql_insert = """ INSERT INTO classification_gz (text_id,
                                                        title, 
                                                        content, 
                                                        site, 
                                                        url, 
                                                        source, 
                                                        publish_time, 
                                                        news_author, 
                                                        mainNavigation, 
                                                        subNavigation, 
                                                        collect_time, 
                                                        pred_label_gz, 
                                                        pred_label_area) 
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ;"""

        try:
            self.cursor.execute(sql_insert, line)
            self.conn.commit()
        except Exception as e:
            self.conn.rollback()
            print("update scores 2 mysql connection exception: %s %s" % (sql_insert, e))

    def connect2db(self, userName, password, dbName, ip="localhost", port=3306):
        try:
            db_coon = pymysql.connect(host=ip, user=userName, password=password, db=dbName, port=port, charset='utf8')
        except Exception as e:
            print('Error: connect mysql falied!')
        else:
            cursor = db_coon.cursor()
            print('Info: connect mysql success!')
        return db_coon,cursor

            

def main():
    try:
        #c = Consumer(KafkaServerList=['10.10.150.43:9092'], Topics=['mysqldbserver1.jkw.js_news'])
        p = Producer(KafkaServerList=['10.10.149.20:6667'], ClientId='kuankuan', Topic='junshi-guanzhu-six')
        c = Consumer(KafkaServerList=['10.10.149.20:6667'], ClientId='kuankuan', Topics=['junshi-guanzhu', ])
        c.consumeMsg(p)
        p.closeConnection()
    except Exception as err:
        print(err)


if __name__ == "__main__":
    try:
        main()
    finally:
        sys.exit()

